# Copyright 2021-2023 @ Shenzhen Bay Laboratory &
#                       Peking University &
#                       Huawei Technologies Co., Ltd
#
# This code is a part of MindSPONGE:
# MindSpore Simulation Package tOwards Next Generation molecular modelling.
#
# MindSPONGE is open-source software based on the AI-framework:
# MindSpore (https://www.mindspore.cn/)
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""Metadynamics"""

from typing import Union

import torch
from torch import Tensor

from ...potential.bias import Bias
from ...colvar import Colvar
from ...system import Molecule
from ...function import get_integer, get_tensor


class GaussianKernel(Bias):
    r"""Gaussian type kernel"""

    def __init__(self,
                 num_walker: int,
                 colvar: Colvar = None,
                 update_pace: int = 20,
                 bandwidth: Union[float, Tensor] = 0.1,
                 kernel_min: float = 2,
                 kernel_max: float = 4,
                 kernel_scale: float = 0.1,
                 ):

        super().__init__(
            name='gaussian_kernel',
            colvar=colvar,
            update_pace=update_pace,
        )

        self.num_walker = num_walker

        self.kernel_min = get_tensor(kernel_min, torch.float32, self.device)
        self.kernel_max = get_tensor(kernel_max, torch.float32, self.device)
        self.kernel_mid = (self.kernel_max + self.kernel_min) / 2

        kernel_scale = get_tensor(kernel_scale, torch.float32, self.device)
        self.kernel_scale = get_tensor(kernel_scale, torch.float32, self.device)

        bandwidth = get_tensor(bandwidth, torch.float32, self.device)
        self.bandwidth = get_tensor(bandwidth, torch.float32, self.device)

        # (B, B)
        self.gaussian_exp = torch.zeros(self.num_walker, self.num_walker, dtype=torch.float32, device=self.device)

        self.eyes = torch.eye(self.num_walker, self.num_walker, dtype=torch.bool, device=self.device)

        self.m_min = self.kernel_min / (self.num_walker - 1)
        self.m_max = self.kernel_max / (self.num_walker - 1)

    def update(self, coordinate: Tensor, pbc_box: Tensor = None) -> Tensor:
        r"""update bandwidth"""
        bandwidth = self.get_bandwidth(self.gaussian_exp, self.bandwidth)
        self.bandwidth = bandwidth
        return bandwidth

    def initialize(self, system: Molecule) -> Tensor:
        r"""initialize bandwidth"""
        coordinate = system.coordinate
        pbc_box = system.pbc_box
        gaussian_exp = self.calc_gauss_exp(coordinate, pbc_box)
        bandwidth = self.get_bandwidth(gaussian_exp)
        self.bandwidth = bandwidth
        return bandwidth

    def get_bandwidth(self, gaussian_exp: Tensor, bandwidth: Tensor = None) -> Tensor:
        r"""get optimized bandwidth"""
        def _calc_mean(gauss_exp, bandwidth_):
            k = torch.exp(gauss_exp /(torch.square(bandwidth_)))
            return torch.mean(torch.sum(k, dim=-1))

        if bandwidth is not None:
            k_mean = _calc_mean(gaussian_exp, bandwidth)
            if self.kernel_min < k_mean < self.kernel_max:
                return bandwidth

        k_min = torch.amin(gaussian_exp)
        k_max = torch.amax(torch.where(self.eyes, k_min, gaussian_exp))

        s_min = torch.sqrt(k_max / torch.log(self.m_min))
        s_max = torch.sqrt(k_min / torch.log(self.m_max))

        if bandwidth is not None:
            s_max = torch.where(torch.logical_and(k_mean > self.kernel_max, s_max > bandwidth), bandwidth, s_max)
            s_min = torch.where(torch.logical_and(k_mean < self.kernel_min, s_min < bandwidth), bandwidth, s_min)

        bandwidth = (s_min + s_max) / 2
        k_mean = _calc_mean(gaussian_exp, bandwidth)

        while k_mean > self.kernel_max or k_mean < self.kernel_min:
            s_max = torch.where(k_mean > self.kernel_max, bandwidth, s_max)
            s_min = torch.where(k_mean < self.kernel_min, bandwidth, s_min)

            bandwidth = (s_min + s_max) / 2
            k_mean = _calc_mean(gaussian_exp, bandwidth)

        return bandwidth

    def calc_gauss_exp(self, coordinate, pbc_box):
        r"""calcaulte gaussian exponent"""
        mw = coordinate.shape[0]

        colvar = coordinate
        if self.colvar is not None:
            # (B, s_1, s_2, ..., s_n)
            colvar = self.colvar(coordinate, pbc_box)

        colvar_sg = colvar.detach()
        # (B, S) <- (B, s_1, s_2, ..., s_n)
        cv_a = colvar.view(mw, -1)
        cv_b = colvar_sg.view(mw, -1)

        dim = cv_a.shape[-1]

        # (S, B) <- (B, S)
        cv_b = cv_b.t()

        # (B, 1) <- (B, S)
        a2 = torch.sum(cv_a**2, dim=-1, keepdim=True)
        # (1, B) <- (S, B)
        b2 = torch.sum(cv_b**2, dim=0, keepdim=True)

        # (B, B) = (B, S) X (S, B)
        ab = torch.matmul(cv_a, cv_b)

        # (B, B) = (B, 1) + (B, B) + (1, B)
        # (s - s_0) ^ 2 = s^2 - 2 * s * s_0 + s_0^2
        dist2 = (a2 - 2 * ab + b2) / dim

        # (B, B)
        return -0.5 * dist2

    def forward(self,
                  coordinate: Tensor,
                  neighbour_index: Tensor = None,
                  neighbour_mask: Tensor = None,
                  neighbour_vector: Tensor = None,
                  neighbour_distance: Tensor = None,
                  pbc_box: Tensor = None
                  ):
        r"""Calculate kernel.

        Args:
            coordinate (Tensor): Tensor of shape `(B, A, D)`. Data type is float.
                Position coordinate of atoms in system.
            neighbour_index (Tensor): Tensor of shape `(B, A, N)`. Data type is int.
                Index of neighbour atoms. Default: None
            neighbour_mask (Tensor): Tensor of shape `(B, A, N)`. Data type is bool.
                Mask for neighbour atoms. Default: None
            neighbour_vector (Tensor): Tensor of shape `(B, A, N)`. Data type is bool.
                Vectors from central atom to neighbouring atoms.
            neighbour_distance (Tensor): Tensor of shape `(B, A, N)`. Data type is float.
                Distance between neigh_shift atoms. Default: None
            pbc_box (Tensor): Tensor of shape `(B, D)`. Data type is float.
                Tensor of PBC box. Default: None

        Returns:
            kernel (Tensor): Tensor of shape `(B, 1)`. Data type is float.

        Symbols:
            B:  Batchsize, i.e. number of walkers in simulation
            A:  Number of atoms.
            N:  Maximum number of neighbour atoms.
            D:  Spatial dimension of the simulation system. Usually is 3.

        """

        # (B, B)
        gauss_exp = self.calc_gauss_exp(coordinate, pbc_box)
        self.gaussian_exp = gauss_exp

        # (B, B)
        kernel_exp = gauss_exp * torch.reciprocal(self.bandwidth**2)
        # print(F.reduce_mean(kernel_exp))

        # (B, 1) <- (B, B)
        kernel = torch.logsumexp(kernel_exp, dim=-1, keepdim=True)
        # kernel = self.logsumexp(kernel_exp)

        return kernel
